import os

import torch
import numpy as np
import scipy.signal

from matplotlib import pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def triplet_loss(alpha=0.2):
    def _triplet_loss(y_predict, batch_size):
        anchor = y_predict[:int(batch_size)]
        positive = y_predict[int(batch_size):int(2 * batch_size)]
        negative = y_predict[int(2 * batch_size):]

        print(anchor.shape)
        print(positive.shape)
        pos_dist = torch.sqrt(torch.sum(torch.pow(anchor - positive, 2), axis=-1))
        neg_dist = torch.sqrt(torch.sum(torch.pow(anchor - negative, 2), axis=-1))

        keep_all = (neg_dist - pos_dist < alpha).cpu().numpy().flatten()
        hard_triplets = np.where(keep_all == 1)

        pos_dist = pos_dist[hard_triplets].to(device)
        neg_dist = neg_dist[hard_triplets].to(device)

        basic_loss = pos_dist - neg_dist + alpha
        loss = torch.sum(basic_loss) / torch.max(torch.tensor(1), torch.tensor(len(hard_triplets[0])))
        return loss

    return _triplet_loss


class LossHistory:
    def __init__(self, log_dir):
        import datetime
        curr_time = datetime.datetime.now()
        time_str = datetime.datetime.strftime(curr_time, "%Y_%m_%d_%H_%M_%S")
        self.log_dir = log_dir
        self.time_str = time_str
        self.save_path = os.path.join(self.log_dir, "loss_" + str(self.time_str))
        self.acc = []
        self.losses = []
        self.val_loss = []

        os.makedirs(self.save_path)

    def append_loss(self, acc, loss, val_loss):
        self.acc.append(acc)
        self.losses.append(loss)
        self.val_loss.append(val_loss)
        with open(os.path.join(self.save_path, "epoch_acc_" + str(self.time_str) + ".txt"), 'a') as f:
            f.write(str(acc))
            f.write("\n")
        with open(os.path.join(self.save_path, "epoch_loss_" + str(self.time_str) + ".txt"), 'a') as f:
            f.write(str(loss))
            f.write("\n")
        with open(os.path.join(self.save_path, "epoch_val_loss_" + str(self.time_str) + ".txt"), 'a') as f:
            f.write(str(val_loss))
            f.write("\n")
        self.loss_plot()

    def loss_plot(self):
        iters = range(len(self.losses))

        plt.figure()
        plt.plot(iters, self.losses, 'red', linewidth=2, label='train loss')
        plt.plot(iters, self.val_loss, 'coral', linewidth=2, label='val loss')
        try:
            if len(self.losses) < 25:
                num = 5
            else:
                num = 15

            tmp = scipy.signal.savgol_filter(self.losses, num, 3)
            plt.plot(iters, tmp, 'green', linestyle='--', linewidth=2, label='smooth train loss')
            tmp = scipy.signal.savgol_filter(self.val_loss, num, 3)
            plt.plot(iters, tmp, '#8B4513', linestyle='--', linewidth=2, label='smooth val loss')
        except Exception as e:
            print(e)
            pass

        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend(loc="upper right")

        plt.savefig(os.path.join(self.save_path, "epoch_loss_" + str(self.time_str) + ".png"))

        plt.cla()
        plt.close("all")

        plt.figure()
        plt.plot(iters, self.acc, 'red', linewidth=2, label='acc')
        try:
            num = 5 if len(self.losses) < 25 else 15
            tmp = scipy.signal.savgol_filter(self.acc, num, 3)
            plt.plot(iters, tmp, 'green', linestyle='--', linewidth=2, label='smooth acc')
        except Exception as e:
            print(e)
            pass

        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Lfw Acc')
        plt.legend(loc="upper right")

        plt.savefig(os.path.join(self.save_path, "epoch_acc_" + str(self.time_str) + ".png"))

        plt.cla()
        plt.close("all")
